# switch_client.py
import paramiko
import re
import time


class HuaweiSwitchClient:
    def __init__(self, host, port, username, password):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.client = None
        self.shell = None

    # ---------- 基础连接 ----------
    def connect(self):
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.client.connect(self.host, port=self.port,
                            username=self.username, password=self.password, timeout=30)
        self.shell = self.client.invoke_shell()
        time.sleep(0.1)

    def close(self):
        if self.shell:
            self.shell.close()
        if self.client:
            self.client.close()

    # ---------- 稳健读取 ----------
    def _read_until(self, expect_list=None, timeout=15):
        """读到指定正则列表任意一项出现为止"""
        if expect_list is None:
            expect_list = [r'<[^>]+>', r'\[.+\]']
        buf = ''
        end_time = time.time() + timeout
        while time.time() < end_time:
            if self.shell.recv_ready():
                chunk = self.shell.recv(4096).decode('utf-8', errors='ignore')
                buf += chunk
                for pattern in expect_list:
                    if re.search(pattern, buf):
                        return buf
            time.sleep(0.1)
        raise TimeoutError('SSH 读取超时')

    def send_cmd(self, cmd, expect_list=None, timeout=15):
        """发送命令并返回完整回显"""
        self.shell.send(cmd + '\n')
        out = self._read_until(expect_list, timeout)
        # 遇到分页就空格翻页
        while re.search(r'---- More ----|--More--', out):
            self.shell.send(' ')
            out += self._read_until(expect_list, timeout)
        return out

    def get_bindings(self, subnet):
        self.connect()
        self.send_cmd("screen-length 0 temporary")
        self.send_cmd("system")
        cfg = self.send_cmd("display current-configuration | include arp static")
        self.send_cmd("quit")
        self.close()

        plain = re.sub(r'\x1b\[[0-9;]*m', '', cfg)
        # ★★ 只抓 IP、MAC、接口（可选） ★★
        static_pattern = re.compile(
            r'arp\s+static\s+'
            r'(?P<ip>\d+\.\d+\.\d+\.\d+)\s+'
            r'(?P<mac>[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4})'
            r'(?:\s+vid\s+\d+)?'
            r'(?:\s+interface\s+(?P<interface>\S+))?',
            re.I
        )

        bindings = []
        for m in static_pattern.finditer(plain):
            ip = m.group('ip')
            mac = m.group('mac')
            iface = m.group('interface') or 'N/A'
            if subnet and subnet != "all" and not ip.startswith(subnet + '.'):
                continue
            bindings.append({
                'ip_address': ip,
                'mac_address': mac,
                'interface': iface
            })

        bindings.sort(key=lambda x: int(x['ip_address'].split('.')[-1]))
        return bindings
